#pragma once
#include <cassert>
#include <algorithm>
#include <iostream>

namespace STL
{
    enum color
    {
        RED = 0,
        BLACK
    };

    template <class Val>
    class RBNode
    {
    public:
        Val _data;
        RBNode *_parent;
        RBNode *_leftchild;
        RBNode *_rightchild;
        color _color;

        RBNode(Val data)
            : _data(data), _parent(nullptr), _leftchild(nullptr), _rightchild(nullptr), _color(RED)
        {
        }

        RBNode(const RBNode &tmp)
            : _data(tmp._data), _parent(tmp._parent), _leftchild(tmp._leftchild), _rightchild(tmp._rightchild), _color(tmp._color)
        {
        }
    };

    template <class T, class Pointer, class Ref>
    class Iterator
    {
        using Node = RBNode<T>;
        using Self = Iterator<T, Pointer, Ref>;

    public:
        Iterator(Node *node, Node *root) : _node(node), _root(root) {}

        Self &operator++()
        {
            assert(_node != nullptr);
            if (_node->_rightchild != nullptr)
            {
                // 1、当前分支没有遍历完，寻找后继
                Node *tmp = _node->_rightchild;
                while (tmp->_leftchild)
                    tmp = tmp->_leftchild;
                _node = tmp;
            }
            else
            {
                // 2、当前分支遍历完，寻找第一个右边的祖先节点
                Node *tmp = _node;
                Node *parent = tmp->_parent;
                while (parent && parent->_rightchild == tmp)
                {
                    tmp = parent;
                    parent = tmp->_parent;
                }
                _node = parent;
            }
            return *this;
        }

        Self &operator--()
        {
            if (_node == nullptr)
            {
                // 1、当前节点在end()
                Node *tmp = _root;
                while (tmp->_rightchild)
                {
                    tmp = tmp->_rightchild;
                }
                _node = tmp;
            }
            else if (_node->_leftchild != nullptr)
            {
                // 2、当前分支没有遍历完，当前结点前驱
                Node *tmp = _node->_leftchild;
                while (tmp->_rightchild)
                    tmp = tmp->_rightchild;
                _node = tmp;
            }
            else
            {
                // 3、当前分支遍历完，第一个左边的祖先结点
                Node *tmp = _node;
                Node *parent = _node->_parent;
                while (parent && parent->_leftchild == tmp)
                {
                    tmp = parent;
                    parent = parent->_parent;
                }
                _node = parent;
            }
            return *this;
        }

        Ref operator*()
        {
            return _node->_data;
        }

        Pointer operator->()
        {
            return &_node->_data;
        }

        bool operator!=(const Self &s) const
        {
            return _node != s._node;
        }

        bool operator==(const Self &s) const
        {
            return _node == s._node;
        }

    private:
        Node *_node;
        Node *_root;
    };

    template <class Key, class Val, class KeyOfT, class Compare>
    class RBTree
    {
        using TreeNode = RBNode<Val>;

    public:
        using iterator = Iterator<Val, Val *, Val &>;
        using const_iterator = Iterator<Val, const Val *, const Val &>;

    public:
        RBTree() = default;

        RBTree(const RBTree &tmp)
        {
            _root = copy(tmp._root);
        }

        ~RBTree()
        {
            destroy();
        }

        iterator begin()
        {
            if (_root == nullptr)
                return iterator(nullptr, nullptr);
            TreeNode *tmp = _root;
            while (tmp->_leftchild)
                tmp = tmp->_leftchild;
            return iterator(tmp, _root);
        }

        iterator end()
        {
            return iterator(nullptr, _root);
        }

        const_iterator begin() const
        {
            if (_root == nullptr)
                return const_iterator(nullptr, nullptr);
            TreeNode *tmp = _root;
            while (tmp->_leftchild)
                tmp = tmp->_leftchild;
            return const_iterator(tmp, _root);
        }

        const_iterator end() const
        {
            return const_iterator(nullptr, _root);
        }

        bool insert(const Val &in)
        {
            // 1、找到合适插入的结点位置
            TreeNode *parent = nullptr;
            TreeNode *cur = _root;
            KeyOfT kot;
            Compare comp;
            while (cur)
            {
                TreeNode *lchild = cur->_leftchild;
                TreeNode *rchild = cur->_rightchild;
                if (comp(kot(cur->_data), kot(in)))
                {
                    parent = cur;
                    cur = rchild;
                }
                else if (comp(kot(in), kot(cur->_data)))
                {
                    parent = cur;
                    cur = lchild;
                }
                else
                {
                    return false;
                }
            }

            if (parent == nullptr)
            {
                _root = new TreeNode(in);
                _root->_color = BLACK;
                return true;
            }
            cur = new TreeNode(in);
            cur->_parent = parent;
            if (comp(kot(parent->_data), kot(in)))
            {
                // 插入右边
                parent->_rightchild = cur;
            }
            else
            {
                // 插入左边
                parent->_leftchild = cur;
            }

            // 进行颜色调整/旋转
            TreeNode *grandparent = nullptr;
            TreeNode *uncle = nullptr;
            // 当grandparent == _root时，旋转要更改指向
            while (parent && parent->_color == RED)
            {
                grandparent = parent->_parent;
                if (parent == grandparent->_leftchild)
                {
                    uncle = grandparent->_rightchild;
                    if (uncle && uncle->_color == RED)
                    {
                        // 变色然后循环向上更新
                        parent->_color = BLACK;
                        uncle->_color = BLACK;
                        grandparent->_color = RED;
                        cur = grandparent;
                        parent = cur->_parent;
                    }
                    else
                    {
                        if (cur == parent->_leftchild)
                        {
                            rotateR(grandparent);
                            parent->_color = BLACK;
                            grandparent->_color = RED;
                            break;
                        }
                        else if (cur == parent->_rightchild)
                        {
                            rotateLR(grandparent);
                            cur->_color = BLACK;
                            grandparent->_color = RED;
                            break;
                        }
                        else
                        {
                            assert(false);
                        }
                    }
                }
                else if (parent == grandparent->_rightchild)
                {
                    uncle = grandparent->_leftchild;
                    if (uncle && uncle->_color == RED)
                    {
                        // 变色然后循环向上更新
                        parent->_color = BLACK;
                        uncle->_color = BLACK;
                        grandparent->_color = RED;
                        cur = grandparent;
                        parent = cur->_parent;
                    }
                    else
                    {
                        if (cur == parent->_rightchild)
                        {
                            rotateL(grandparent);
                            parent->_color = BLACK;
                            grandparent->_color = RED;
                            break;
                        }
                        else if (cur == parent->_leftchild)
                        {
                            rotateRL(grandparent);
                            cur->_color = BLACK;
                            grandparent->_color = RED;
                            break;
                        }
                        else
                        {
                            assert(false);
                        }
                    }
                }
                else
                {
                    assert(false);
                }
            }
            _root->_color = BLACK;
            return true;
        }

        void inorder() const
        {
            std::cout << "中序：";
            _inorder(_root);
            std::cout << "\n";
        }

        // void _inorder(TreeNode *node) const
        // {
        //     if (!node)
        //         return;
        //     _inorder(node->_leftchild);
        //     std::cout << "{" << node->_data << "} ";
        //     _inorder(node->_rightchild);
        // }

        void _inorder(TreeNode *node) const
        {
            if (!node)
                return;
            _inorder(node->_leftchild);
            std::cout << "{" << node->_data.first << "," << node->_data.second << "} ";
            _inorder(node->_rightchild);
        }

        bool check_redblack_tree()
        {
            TreeNode *p = _root;
            size_t reference = 0;
            while (p)
            {
                if (p->_color == BLACK)
                    reference += 1;
                p = p->_leftchild;
            }
            return _check_redblack_tree(_root, 0, reference);
        }

        void destroy()
        {
            _destroy(_root);
        }

        TreeNode *copy(TreeNode *node) const
        {
            if (node == nullptr)
                return nullptr;
            TreeNode *newnode = new TreeNode(*node);
            newnode->_leftchild = copy(node->_leftchild);
            newnode->_rightchild = copy(node->_rightchild);
            return newnode;
        }

    private:
        void _destroy(TreeNode *node)
        {
            if (node == nullptr)
                return;
            _destroy(node->_leftchild);
            _destroy(node->_rightchild);
            delete (node);
        }

        void rotateL(TreeNode *parent)
        {
            TreeNode *grandparent = parent->_parent;
            TreeNode *rchild = parent->_rightchild;
            TreeNode *rlchild = rchild->_leftchild;
            // 1、锁定rlchild
            parent->_rightchild = rlchild;
            if (rlchild)
                rlchild->_parent = parent;
            // 2、锁定parent
            rchild->_leftchild = parent;
            parent->_parent = rchild;
            // 3、锁定rchild
            if (grandparent == nullptr)
            {
                rchild->_parent = nullptr;
                _root = rchild;
            }
            else
            {
                if (grandparent->_rightchild == parent)
                {
                    grandparent->_rightchild = rchild;
                }
                else if (grandparent->_leftchild == parent)
                {
                    grandparent->_leftchild = rchild;
                }
                else
                {
                    assert(false);
                }
                rchild->_parent = grandparent;
            }
        }

        void rotateR(TreeNode *parent)
        {
            TreeNode *grandparent = parent->_parent;
            TreeNode *lchild = parent->_leftchild;
            TreeNode *lrchild = lchild->_rightchild;
            // 1、锁定lrchild
            parent->_leftchild = lrchild;
            if (lrchild)
                lrchild->_parent = parent;
            // 2、锁定parent
            lchild->_rightchild = parent;
            parent->_parent = lchild;
            // 3、锁定lchild
            if (grandparent == nullptr)
            {
                lchild->_parent = nullptr;
                _root = lchild;
            }
            else
            {
                if (grandparent->_rightchild == parent)
                {
                    grandparent->_rightchild = lchild;
                }
                else if (grandparent->_leftchild == parent)
                {
                    grandparent->_leftchild = lchild;
                }
                else
                {
                    assert(false);
                }
                lchild->_parent = grandparent;
            }
        }

        void rotateLR(TreeNode *parent)
        {
            rotateL(parent->_leftchild);
            rotateR(parent);
        }

        void rotateRL(TreeNode *parent)
        {
            rotateR(parent->_rightchild);
            rotateL(parent);
        }

        bool _check_redblack_tree(TreeNode *node, size_t cur_nums, const size_t &reference)
        {
            if (node == nullptr)
            {
                if (cur_nums != reference)
                    return false;
                else
                    return true;
            }
            if (node->_color == BLACK)
                cur_nums += 1;
            return _check_redblack_tree(node->_leftchild, cur_nums, reference) &&
                   _check_redblack_tree(node->_rightchild, cur_nums, reference);
        }

    private:
        TreeNode *_root = nullptr;
    };
}